# (c) 2025-26 Haodong, Qi; Stefano M., Iacus
#
# ------------------------------------
# Config
# ------------------------------------
# Make sure you have 'pacman' package installed first.
# You can install the package with:  install.packges("pacman")

# If running in RStudio, set working dir to script location (optional):
# uncomment the lines below and execute them
#
# if (requireNamespace("rstudioapi", quietly = TRUE) &&
#   rstudioapi::isAvailable()) {
#     setwd(dirname(rstudioapi::getActiveDocumentContext()$path))
#   }

pacman::p_load(lavaan, semPlot, data.table, tidyverse, ggcorrplot,
  ComplexHeatmap, circlize)

# ------------------------------------
# data
# ------------------------------------

df <- fread(file.path(".", "data_avg.csv"))
domain <- fread(file.path(".","variable_domain_map.csv"))
setorder(domain, domain)

# 1) Add a row id to your original data
df[, id := .I]

# reverse negativity for better interpretation 
df[,  fin_sec:=finworry* -1]
df[,  future_sec:=fearfuture* -1]
df[, !c("finworry", "fearfuture"), with = FALSE]

# ------------------------------------
# SEM Analysis
# ------------------------------------

# Specify model
model <- '
  # -- Measurement model
  clim_risk =~ HEAT_RISK + FIRE_RISK + DROUGHT_RISK +
    COASTAL_RISK + WIND_RISK + INLAND_RISK 

  charac_virt =~ altruism + charity + delaygrat + empathy + forgive +
    goodpromo + gratitude + volunteer
  relig =~ afterlife + believegod + lovedgod + relcomfort + relcrit + spiritpun
  inst_trust =~ govapprove + polvoice + discrim
  meaning =~ balance + hope + purpose + resilience
  social =~ belonging + loneliness + relationships + trust + trusted 
  subj_welb =~ happiness + innerpeace + lifesat + optimism
  psy_dstrs =~ anxiety + depression + suffering + selfesteem
  physc_cond =~ energy + posfunct + vitality + healthlim + pain 
  econ_stab =~ jobsat + mastery + fin_sec + future_sec

  # -- Structural model
  inst_trust ~ clim_risk  
  meaning ~ clim_risk 
  social ~ clim_risk  
  subj_welb ~ clim_risk  
  psy_dstrs ~ clim_risk   
  physc_cond ~ clim_risk  
  econ_stab ~ clim_risk   
  charac_virt ~ clim_risk  
  relig ~ clim_risk

  # -- remove covariance
  inst_trust ~~ 0*meaning + 0*social + 0*subj_welb + 0*psy_dstrs +
    0*physc_cond + 0*econ_stab + 0*charac_virt + 0*relig
  meaning ~~  0*social + 0*subj_welb + 0*psy_dstrs +
    0*physc_cond + 0*econ_stab + 0*charac_virt + 0*relig
  social ~~ 0*subj_welb + 0*psy_dstrs +
    0*physc_cond + 0*econ_stab + 0*charac_virt + 0*relig
  subj_welb ~~ 0*psy_dstrs +
    0*physc_cond + 0*econ_stab + 0*charac_virt + 0*relig
  psy_dstrs ~~ 0*physc_cond + 0*econ_stab + 0*charac_virt + 0*relig
  physc_cond ~~ 0*econ_stab + 0*charac_virt + 0*relig
  econ_stab ~~ 0*charac_virt + 0*relig
  charac_virt ~~ 0*relig

'
# Fit model
fit <- sem(model, data = df, std.lv = TRUE)
summary(fit, standardized = TRUE, fit.measures = TRUE)

fitMeasures(fit, c("df", "cfi", "tli", "rmsea", "srmr", "chisq", "pvalue"))

subset(standardizedSolution(fit), op=="=~" | op=="~")

# !!! plot std sempaths with sig. stars
pe <- standardizedSolution(fit)
pe$stars <- ifelse(pe$pvalue < .001, "***",
             ifelse(pe$pvalue < .01, "**",
             ifelse(pe$pvalue < .05, "*", "")))
struct <- pe[pe$op == "~" | pe$op == "=~", ]
struct$label <- paste0(
  sprintf("%.2f", struct$est.std),
  struct$stars
)

png("path_mlt_outcomes.png", width = 8000, height = 8000, res = 3000)
semPaths(fit, 
  what="std", edgeLabels = struct$label,
  edge.label.cex = 0.6, sizeMan2=2, sizeLat2 = 5,
  curvePivot = TRUE, layout = "tree2",
  nCharNodes = 0, rotation = 2, style = "lisrel",
  residuals = FALSE,       # hides residual arrows
  intercepts = FALSE,      # hides intercepts
  covAtResiduals = FALSE, 
  edge.width = .5,         # constant stroke width for all edges
  asize = 2                # (optional) constant arrowhead size
)
dev.off()

#  Get the mapping of cases used by lavaan back to df row ids
case_idx <- lavInspect(fit, "case.idx")   # integer indices into df

# Get latent scores and attach the correct 'id'
latent_scores <- as.data.table(lavPredict(fit))
latent_scores[, id := case_idx]

latent_scores <- merge(df, latent_scores, by = "id", all.x = TRUE)

fwrite(latent_scores, "latent_scores.csv")


# ------------------------------------
# mapping latent and indicator vars
# ------------------------------------
pe_std <- as.data.table(standardizedSolution(fit))
meas_dt <- pe_std[op == "=~",
  .(latent = lhs, indicator = rhs,
    est_std = est.std, se, z, pvalue)
][order(latent, -abs(est_std))]
domain <- data.table(variable=meas_dt$indicator, label=meas_dt$latent)



# ------------------------------------
# Corrplot
# ------------------------------------
vars <- names(df)
cor_data_avg <- df[, 
  !c("id", "StateCounty", 
  grep("_RAJ|OVERALL|MEAN|worry|fear|migmood|corrup|tweets|rural|population",vars,value=T)), 
  with = FALSE]
vars <- names(cor_data_avg)

cor_matrix_avg <- cor(cor_data_avg, 
  use = "pairwise.complete.obs",  method = "spearman")

# Computing ordering within each group
geo_vars <- grep("RISK",vars,value=T)
swb_vars <- grep("RISK",vars,value = T, invert = T)
geo_order_avg <- hclust(as.dist(1 - cor_matrix_avg[geo_vars, geo_vars]))$order
swb_order_avg <- hclust(as.dist(1 - cor_matrix_avg[swb_vars, swb_vars]))$order

# Combining ordered variable names
ordered_vars_avg <- c(geo_vars[geo_order_avg], swb_vars[swb_order_avg])

# Reordering the full matrix
cor_matrix_ordered_avg <- cor_matrix_avg[ordered_vars_avg, ordered_vars_avg]



## Dropping technical variables and map labels
swb_dt <- data.table(variable = swb_vars)
swb_dt <- merge(swb_dt, domain[, .(variable, label)], by = "variable", all.x = TRUE)
swb_dt <- swb_dt[!is.na(label)]

## Conceptual block order (includes 'civic')
label_order <- c("subj_welb","meaning","charac_virt","social",
                 "psy_dstrs", "physc_cond", "econ_stab","relig",
                 "inst_trust")
swb_dt[, label := factor(label, levels = label_order)]

## Force SWB variables to follow conceptual blocks (no within-block clustering)
order_by_label <- swb_dt[order(label, variable), variable]

## Final variable order: GEO first (clustered), then SWB by label order
ordered_vars <- c(geo_vars[geo_order_avg], order_by_label)

## Build upper-triangular matrix (hide lower triangle + diagonal)
stopifnot(all(ordered_vars %in% rownames(cor_matrix_avg)),
          all(ordered_vars %in% colnames(cor_matrix_avg)))
mat <- cor_matrix_avg[ordered_vars, ordered_vars, drop = FALSE]
mat[lower.tri(mat, diag = TRUE)] <- NA   # upper triangle only (no diagonal)

# Capitalize all variable names for display
require(tools)
colnames(mat) <- toTitleCase(tolower(colnames(mat)))
rownames(mat) <- toTitleCase(tolower(rownames(mat)))

## Domain split factor
# GEO vars show as "Climate Risk Factors"
split_vec <- ifelse(ordered_vars %in% swb_dt$variable,
                    as.character(swb_dt$label[match(ordered_vars, swb_dt$variable)]),
                    "Climate Risk Factors")
split_levels <- c("Climate Risk Factors", label_order)
split_fac <- factor(split_vec, levels = split_levels)

## Color palettes
# Diverging correlation colors
corr_cols <- colorRamp2(c(-1, 0, 1), c("#3B4CC0", "#FFFFFF", "#E66100"))

# Domain colors (aligned to split levels)
domain_colors <- c(
  "Climate Risk Factors" = "#5E81AC",  # blue-gray
  subj_welb     = "#F4A261",           # warm orange
  meaning       = "#2A9D8F",           # teal
  charac_virt   = "#E9C46A",           # golden
  social        = "#264653",           # deep slate
  psy_dstrs     = "#D1495B",           # soft crimson
  physc_cond    = "#90BE6D",           # green
  econ_stab     = "#A06CD5",           # violet
  relig         = "#1F487E",           # navy
  inst_trust    = "#B8B8B8"            # neutral gray
)
domain_colors <- domain_colors[levels(split_fac)]

# Extended legend labels
domain_labels <- c(
  "Climate Risk Factors" = "Climate Risk Factors",
  subj_welb     =  "Subjective Wellbeing",
  meaning       =  "Meaning & Purpose",
  charac_virt   =  "Character Virtue",
  social        = "Social Connectedness",          
  psy_dstrs     = "Psychological Distress",          
  physc_cond    = "Physical Condition",        
  econ_stab     = "Economic Stability",         
  relig         = "Religious & Spiritual",         
  inst_trust    = "Institutional Trust"         
)

## Top + right annotation bars — disable their auto legends
top_annot <- HeatmapAnnotation(
  Domain = split_fac,
  col = list(Domain = domain_colors),
  show_annotation_name = FALSE,
  show_legend = FALSE     # <- IMPORTANT: no default "Domain" legend
)
right_annot <- rowAnnotation(
  Domain = split_fac,
  col = list(Domain = domain_colors),
  show_annotation_name = FALSE,
  show_legend = FALSE     # <- IMPORTANT: no default "Domain" legend
)

## Heatmap
ht <- Heatmap(
  mat,
  name = "Correlation",
  col = corr_cols,
  na_col = "white",
  cluster_rows = FALSE, cluster_columns = FALSE,
  show_row_dend = FALSE, show_column_dend = FALSE,
  row_split = split_fac, column_split = split_fac,
  cluster_row_slices = FALSE, cluster_column_slices = FALSE,
  row_gap = unit(1.5, "mm"), column_gap = unit(1.5, "mm"),
  ## ---- labels
  show_row_names = TRUE,
  show_column_names = TRUE,          # show top labels
  column_names_side = "top",         # put them on top
  column_names_rot  = 45,            # rotate 45°
  column_names_gp   = gpar(fontsize = 7, just = "right"),
  column_names_max_height = unit(8, "lines"),  # give room so rotation is kept
  row_names_side = "right",
  row_names_gp = gpar(fontsize = 7),
  border = NA,
  top_annotation = top_annot,
  right_annotation = right_annot,
  heatmap_legend_param = list(title = "Correlation"),
  row_title = NULL, column_title = NULL
)

pdf("conceptual_domains_heatmap.pdf", width = 12, height = 9, useDingbats = FALSE)
draw(
  ht,
  merge_legend = TRUE,
  heatmap_legend_side = "right",
  annotation_legend_side = "right",
  show_annotation_legend = FALSE,
  annotation_legend_list = list(
    Legend(
      labels = domain_labels,
      legend_gp = gpar(fill = domain_colors),
      title = "Conceptual Domains",
      title_gp = gpar(fontface = "bold"),
      labels_gp = gpar(fontsize = 8),
      row_gap = unit(2, "mm")
    )
  )
)
dev.off()


